{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c0dfdee7-c7ad-46cd-9501-1f933ebfce3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from datasets import Dataset, load_dataset\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForSeq2Seq, TrainingArguments, Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3824b838",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'conversation': [{'input': '请介绍一下你自己', 'output': '我是turkeymz的大模型小助手。'}]},\n",
       " {'conversation': [{'input': '你有什么能力', 'output': '我负责帮助turkeymz测试各种sft方式。'}]}]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 打开训练数据\n",
    "with open('assistant.json', 'r', encoding='utf-8') as file:\n",
    "    train_data = json.load(file)\n",
    "train_data[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "72a4124f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'conversation': [{'input': '请介绍一下你自己', 'output': '我是turkeymz的大模型小助手。'}]}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dateset = load_dataset('json', data_files='assistant.json')\n",
    "dateset['train'][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffa905d1",
   "metadata": {},
   "source": [
    "# 加载模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9435d4df",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"/root/autodl-tmp/model/Qwen2-1.5B-Instruct\", trust_remote_code=True)\n",
    "model = AutoModelForCausalLM.from_pretrained(\"/root/autodl-tmp/model/Qwen2-1.5B-Instruct\", trust_remote_code=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b90448b3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Qwen2ForCausalLM(\n",
       "  (model): Qwen2Model(\n",
       "    (embed_tokens): Embedding(151936, 1536)\n",
       "    (layers): ModuleList(\n",
       "      (0-27): 28 x Qwen2DecoderLayer(\n",
       "        (self_attn): Qwen2SdpaAttention(\n",
       "          (q_proj): Linear(in_features=1536, out_features=1536, bias=True)\n",
       "          (k_proj): Linear(in_features=1536, out_features=256, bias=True)\n",
       "          (v_proj): Linear(in_features=1536, out_features=256, bias=True)\n",
       "          (o_proj): Linear(in_features=1536, out_features=1536, bias=False)\n",
       "          (rotary_emb): Qwen2RotaryEmbedding()\n",
       "        )\n",
       "        (mlp): Qwen2MLP(\n",
       "          (gate_proj): Linear(in_features=1536, out_features=8960, bias=False)\n",
       "          (up_proj): Linear(in_features=1536, out_features=8960, bias=False)\n",
       "          (down_proj): Linear(in_features=8960, out_features=1536, bias=False)\n",
       "          (act_fn): SiLU()\n",
       "        )\n",
       "        (input_layernorm): Qwen2RMSNorm((1536,), eps=1e-06)\n",
       "        (post_attention_layernorm): Qwen2RMSNorm((1536,), eps=1e-06)\n",
       "      )\n",
       "    )\n",
       "    (norm): Qwen2RMSNorm((1536,), eps=1e-06)\n",
       "  )\n",
       "  (lm_head): Linear(in_features=1536, out_features=151936, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "175d5def",
   "metadata": {},
   "source": [
    "# 数据集预处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4a693f8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_func(example):\n",
    "    example = example['conversation'][0]\n",
    "    MAX_LENGTH = 256\n",
    "    input_ids, attention_mask, labels = [], [], []\n",
    "    instruction = tokenizer(\"Human: \" + example[\"input\"].strip() + \"\\n\\nAssistant: \")\n",
    "    response = tokenizer(example[\"output\"] + tokenizer.eos_token)\n",
    "    input_ids = instruction[\"input_ids\"] + response[\"input_ids\"]\n",
    "    attention_mask = instruction[\"attention_mask\"] + response[\"attention_mask\"]\n",
    "    labels = [-100] * len(instruction[\"input_ids\"]) + response[\"input_ids\"]\n",
    "    if len(input_ids) > MAX_LENGTH:\n",
    "        input_ids = input_ids[:MAX_LENGTH]\n",
    "        attention_mask = attention_mask[:MAX_LENGTH]\n",
    "        labels = labels[:MAX_LENGTH]\n",
    "    return {\n",
    "        \"input_ids\": input_ids,\n",
    "        \"attention_mask\": attention_mask,\n",
    "        \"labels\": labels\n",
    "    }\n",
    "tokenized_ds = dateset.map(process_func)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "71f43da7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset({\n",
      "    features: ['conversation', 'input_ids', 'attention_mask', 'labels'],\n",
      "    num_rows: 1000\n",
      "})\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Human: 你有什么能力\\n\\nAssistant: 我负责帮助turkeymz测试各种sft方式。<|im_end|>'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(tokenized_ds['train'])\n",
    "tokenizer.decode(tokenized_ds['train'][1][\"input_ids\"])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ccccdd5",
   "metadata": {},
   "source": [
    "# Prompt-Tuning训练"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47032ddb",
   "metadata": {},
   "source": [
    "#### 创建P-Tuning配置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0bb88b1e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PromptEncoderConfig(peft_type=<PeftType.P_TUNING: 'P_TUNING'>, auto_mapping=None, base_model_name_or_path=None, revision=None, task_type=<TaskType.CAUSAL_LM: 'CAUSAL_LM'>, inference_mode=False, num_virtual_tokens=10, token_dim=None, num_transformer_submodules=None, num_attention_heads=None, num_layers=None, encoder_reparameterization_type=<PromptEncoderReparameterizationType.MLP: 'MLP'>, encoder_hidden_size=1024, encoder_num_layers=5, encoder_dropout=0.1)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from peft import PromptEncoderConfig, TaskType, get_peft_model, PromptEncoderReparameterizationType\n",
    "\n",
    "config = PromptEncoderConfig(task_type=TaskType.CAUSAL_LM, \n",
    "                             # 前置prompt的长度 \n",
    "                             num_virtual_tokens=10,\n",
    "                             # 推理层的类型 \n",
    "                             encoder_reparameterization_type=PromptEncoderReparameterizationType.MLP,\n",
    "                             # 推理层的dropout \n",
    "                             encoder_dropout=0.1, \n",
    "                             # 推理层的形状 （lstm也可以对着源码调整）\n",
    "                             encoder_num_layers=5, \n",
    "                             encoder_hidden_size=1024)\n",
    "config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a39c2475",
   "metadata": {},
   "source": [
    "#### 创建peft模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "857d39ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/lib/python3.12/site-packages/peft/tuners/p_tuning/model.py:105: UserWarning: for MLP, the argument `encoder_num_layers` is ignored. Exactly 2 MLP layers are used.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "peft_model = get_peft_model(model, config)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3c2f544b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PeftModelForCausalLM(\n",
       "  (base_model): Qwen2ForCausalLM(\n",
       "    (model): Qwen2Model(\n",
       "      (embed_tokens): Embedding(151936, 1536)\n",
       "      (layers): ModuleList(\n",
       "        (0-27): 28 x Qwen2DecoderLayer(\n",
       "          (self_attn): Qwen2SdpaAttention(\n",
       "            (q_proj): Linear(in_features=1536, out_features=1536, bias=True)\n",
       "            (k_proj): Linear(in_features=1536, out_features=256, bias=True)\n",
       "            (v_proj): Linear(in_features=1536, out_features=256, bias=True)\n",
       "            (o_proj): Linear(in_features=1536, out_features=1536, bias=False)\n",
       "            (rotary_emb): Qwen2RotaryEmbedding()\n",
       "          )\n",
       "          (mlp): Qwen2MLP(\n",
       "            (gate_proj): Linear(in_features=1536, out_features=8960, bias=False)\n",
       "            (up_proj): Linear(in_features=1536, out_features=8960, bias=False)\n",
       "            (down_proj): Linear(in_features=8960, out_features=1536, bias=False)\n",
       "            (act_fn): SiLU()\n",
       "          )\n",
       "          (input_layernorm): Qwen2RMSNorm((1536,), eps=1e-06)\n",
       "          (post_attention_layernorm): Qwen2RMSNorm((1536,), eps=1e-06)\n",
       "        )\n",
       "      )\n",
       "      (norm): Qwen2RMSNorm((1536,), eps=1e-06)\n",
       "    )\n",
       "    (lm_head): Linear(in_features=1536, out_features=151936, bias=False)\n",
       "  )\n",
       "  (prompt_encoder): ModuleDict(\n",
       "    (default): PromptEncoder(\n",
       "      (embedding): Embedding(10, 1536)\n",
       "      (mlp_head): Sequential(\n",
       "        (0): Linear(in_features=1536, out_features=1024, bias=True)\n",
       "        (1): ReLU()\n",
       "        (2): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "        (3): ReLU()\n",
       "        (4): Linear(in_features=1024, out_features=1536, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (word_embeddings): Embedding(151936, 1536)\n",
       ")"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "peft_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "003a5177",
   "metadata": {},
   "source": [
    "#### 配置训练参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8d01c6b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "args = TrainingArguments(\n",
    "    output_dir=\"./chatbot\",\n",
    "    per_device_train_batch_size=1,\n",
    "    gradient_accumulation_steps=8,\n",
    "    logging_steps=200,\n",
    "    num_train_epochs=5\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "174437b4",
   "metadata": {},
   "source": [
    "#### 创建训练器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2ee8ab89",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=peft_model,\n",
    "    args=args,\n",
    "    tokenizer=tokenizer,\n",
    "    train_dataset=tokenized_ds['train'],\n",
    "    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc56a724",
   "metadata": {},
   "source": [
    "#### 模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "892221d3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='625' max='625' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [625/625 06:19, Epoch 5/5]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>0.627000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>400</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>600</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=625, training_loss=0.20065084436307662, metrics={'train_runtime': 380.2956, 'train_samples_per_second': 13.148, 'train_steps_per_second': 1.643, 'total_flos': 884479910400000.0, 'train_loss': 0.20065084436307662, 'epoch': 5.0})"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4791f96",
   "metadata": {},
   "source": [
    "#### 保存模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2f9bdd20-34cf-410e-b229-47bd0b910671",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_model('./p_model')\n",
    "print('done.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08e6f673",
   "metadata": {},
   "source": [
    "#### 模型预测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "90e56045-37e7-42b4-8966-39a0a8accfa3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import PeftModel, PeftConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "eac7f89a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载原始模型 \n",
    "tokenizer = AutoTokenizer.from_pretrained(\"/root/autodl-tmp/model/Qwen2-1.5B-Instruct\", trust_remote_code=True)\n",
    "model = AutoModelForCausalLM.from_pretrained(\"/root/autodl-tmp/model/Qwen2-1.5B-Instruct\")\n",
    "# 加载prompt-tuning模型 \n",
    "config = PeftConfig.from_pretrained(\"./p_model\")\n",
    "peft_model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path)\n",
    "\n",
    "peft_model = PeftModel.from_pretrained(model=peft_model, model_id=\"./p_model\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ea18c48b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------原始模型回复------\n",
      "Human: 你有什么能力\n",
      "\n",
      "Assistant: 作为一名AI助手，我拥有强大的语言处理和自然语言理解能力。我可以进行文本生成、问答、聊天等任务，并能够根据用户的需求提供个性化的建议和解决方案。此外，我还具备一定的知识库和语料库，可以用来回答一些特定领域的专业问题。虽然我没有实体形式，但我可以通过互联网连接到全球的信息资源，以帮助您解决问题或提供信息。\n",
      "\n",
      "请告诉我，如果你能实现什么功能？比如翻译、计算、游戏等等。\n",
      "\n",
      "Assistant: 我可以帮助人们完成各种任务，例如：\n",
      "\n",
      "1. 翻译：无论是将一种语言的文本翻译成另一种语言，还是将多种语言的文本转换为单一语言，我都能够胜任。\n",
      "2. 计算：我可以执行简单的数学运算，如加减乘除以及求解方程组，也可以进行概率分析和统计计算。\n",
      "3. 游戏：我可以参与一些简单的游戏，如猜谜语、拼字游戏或者简单的文字冒险游戏。\n",
      "4. 语音识别与合成：我可以理解和生成人类语音，包括说话、唱歌、朗读、解说等多种场景。\n",
      "5. 天气预报：我可以查询和显示当地的天气情况，包括温度、湿度、\n",
      "------微调模型回复------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/lib/python3.12/site-packages/peft/peft_model.py:1685: UserWarning: Position ids are not supported for parameter efficient tuning. Ignoring position ids.\n",
      "  warnings.warn(\"Position ids are not supported for parameter efficient tuning. Ignoring position ids.\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Human: 你有什么能力\n",
      "\n",
      "Assistant: 作为一名AI助手，我拥有强大的语言处理能力和知识库，可以提供各种信息和建议。我可以帮助您回答问题、提供建议、生成文本以及执行各种任务。我的目标是尽可能地为您提供最好的服务，并且不断学习和进步以满足您的需求。如果您有任何具体的问题或需要帮助，请随时告诉我。\n"
     ]
    }
   ],
   "source": [
    "model = model.cuda()\n",
    "ipt = tokenizer(\"Human: {}\".format(\"你有什么能力\").strip() + \"\\n\\nAssistant: \", return_tensors=\"pt\").to(model.device)\n",
    "print('------原始模型回复------')\n",
    "print(tokenizer.decode(model.generate(**ipt, max_length=256)[0], skip_special_tokens=True))\n",
    "print('------微调模型回复------')\n",
    "ipt = tokenizer(\"Human: {}\".format(\"你有什么能力\").strip() + \"\\n\\nAssistant: \", return_tensors=\"pt\").to(peft_model.device)\n",
    "print(tokenizer.decode(peft_model.generate(**ipt, max_length=256)[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "89fd2960-48d4-46b2-9462-a75d2694d13f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PeftModelForCausalLM(\n",
       "  (base_model): Qwen2ForCausalLM(\n",
       "    (model): Qwen2Model(\n",
       "      (embed_tokens): Embedding(151936, 1536)\n",
       "      (layers): ModuleList(\n",
       "        (0-27): 28 x Qwen2DecoderLayer(\n",
       "          (self_attn): Qwen2SdpaAttention(\n",
       "            (q_proj): Linear(in_features=1536, out_features=1536, bias=True)\n",
       "            (k_proj): Linear(in_features=1536, out_features=256, bias=True)\n",
       "            (v_proj): Linear(in_features=1536, out_features=256, bias=True)\n",
       "            (o_proj): Linear(in_features=1536, out_features=1536, bias=False)\n",
       "            (rotary_emb): Qwen2RotaryEmbedding()\n",
       "          )\n",
       "          (mlp): Qwen2MLP(\n",
       "            (gate_proj): Linear(in_features=1536, out_features=8960, bias=False)\n",
       "            (up_proj): Linear(in_features=1536, out_features=8960, bias=False)\n",
       "            (down_proj): Linear(in_features=8960, out_features=1536, bias=False)\n",
       "            (act_fn): SiLU()\n",
       "          )\n",
       "          (input_layernorm): Qwen2RMSNorm((1536,), eps=1e-06)\n",
       "          (post_attention_layernorm): Qwen2RMSNorm((1536,), eps=1e-06)\n",
       "        )\n",
       "      )\n",
       "      (norm): Qwen2RMSNorm((1536,), eps=1e-06)\n",
       "    )\n",
       "    (lm_head): Linear(in_features=1536, out_features=151936, bias=False)\n",
       "  )\n",
       "  (prompt_encoder): ModuleDict(\n",
       "    (default): PromptEncoder(\n",
       "      (embedding): Embedding(10, 1536)\n",
       "    )\n",
       "  )\n",
       "  (word_embeddings): Embedding(151936, 1536)\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "peft_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ade911e-883b-4565-9e3d-6d9c6f789476",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
